package opencv.service.algorithms;

import static com.googlecode.javacv.jna.cv.CV_RGB2GRAY;
import static com.googlecode.javacv.jna.cv.cvCalcOpticalFlowPyrLK;
import static com.googlecode.javacv.jna.cv.cvCvtColor;
import static com.googlecode.javacv.jna.cv.cvFindCornerSubPix;
import static com.googlecode.javacv.jna.cv.cvGoodFeaturesToTrack;
import static com.googlecode.javacv.jna.cxcore.CV_TERMCRIT_EPS;
import static com.googlecode.javacv.jna.cxcore.CV_TERMCRIT_ITER;
import static com.googlecode.javacv.jna.cxcore.IPL_DEPTH_32F;
import static com.googlecode.javacv.jna.cxcore.IPL_DEPTH_8U;
import static com.googlecode.javacv.jna.cxcore.cvCopy;
import static com.googlecode.javacv.jna.cxcore.cvCreateImage;
import static com.googlecode.javacv.jna.cxcore.cvGetSize;
import static com.googlecode.javacv.jna.cxcore.cvReleaseImage;
import static com.googlecode.javacv.jna.cxcore.cvSize;
import static com.googlecode.javacv.jna.cxcore.cvTermCriteria;

import java.awt.Color;
import java.awt.Graphics;
import java.awt.Image;
import java.awt.image.BufferedImage;
import java.util.ArrayList;
import java.util.List;

import opencv.service.util.Vector2D;

import com.googlecode.javacv.jna.cxcore.CvPoint2D32f;
import com.googlecode.javacv.jna.cxcore.CvSize;
import com.googlecode.javacv.jna.cxcore.CvTermCriteria;
import com.googlecode.javacv.jna.cxcore.IplImage;

import com.sun.jna.ptr.IntByReference;

/**
 * @author Aakash Polra
 */
public class LucasKanadeOpticalFlow implements MotionDetectionAlgorithm {

	public static final int MAX_CORNERS = 500;
	
	private static double lk_quality = 0.10;	// usually 0.10 or 0.01
	private static double lk_min_distance = 5.0;	// two returned points are not within these number of pixels
	private static CvSize lk_zero_zone = cvSize(-1, -1);	// no zero-zone desired. ( a window whose content should be ignored when
															// taking into account when calculating correlation matrix. (As sometimes it might not result into an invertible matrix)
	private static CvSize lk_window_size = cvSize(10, 10);	// 10 + 1 + 10 = 21 pixels wide and high window
	private static CvTermCriteria lk_termination_criteria = cvTermCriteria(CV_TERMCRIT_ITER|CV_TERMCRIT_EPS, 20, 0.03);	// terminate after 20 iterations or when quality reached to 0.03
	private static int lk_pyr_level = 5;

	
	private boolean debugDrawingEnabled;

	private List<Vector2D> allVectors;
	
	static float threshold = 0.1f;	// TODO find out a good value
	static float minThreshold = 0.3f;	// you have to have this min avg in at least one direction to consider it a z movement
	
	public LucasKanadeOpticalFlow(boolean debugDrawingEnabled) {
		this.debugDrawingEnabled = debugDrawingEnabled;
		this.allVectors = new ArrayList<Vector2D>();
	}
	
	@Override
	public void run(IplImage img1, IplImage img2) {
		Vector2D result = calculateMotion(img1, img2);
		onMotionDetection(result);
		
		if (debugDrawingEnabled) {
			// Prepare debug image with all vectors drawn
			BufferedImage debugImage = img2.getBufferedImage();
			Graphics g = debugImage.getGraphics();
			g.setColor(Color.red);
			for (Vector2D vector : allVectors) {
				g.drawLine((int)vector.getX1(), (int)vector.getY1(), (int)vector.getX2(), (int)vector.getY2());
			}
			onDebugImagePrepared(debugImage);
		}
	}
	
	@Override
	public void onMotionDetection(Vector2D vector) {
		// override this
	}

	@Override
	public void onDebugImagePrepared(Image image) {
		// override this
	}
	
	public Vector2D calculateMotion(IplImage img1Raw, IplImage img2Raw)	{
		
		CvSize imgSize = cvGetSize(img1Raw);

		// Get trackable features
		IplImage eigImg = cvCreateImage(imgSize.byValue(), IPL_DEPTH_32F, 1);
		IplImage tmpImg = cvCreateImage(imgSize.byValue(), IPL_DEPTH_32F, 1);

		IntByReference cornerCount = new IntByReference(MAX_CORNERS);	// this will be set by the func below
		CvPoint2D32f[] corners = CvPoint2D32f.createArray(MAX_CORNERS);	// features. this will be populated by the func below

		if (img1Raw == null || img2Raw == null)
		{
			throw new NullPointerException("null image found in calculateMotion()");
		}

		// Images passed into cvGoodFeaturesToTrack() need to be single channel.
		IplImage img1 = cvCreateImage(cvGetSize(img1Raw), IPL_DEPTH_8U, 1);
		IplImage img2 = cvCreateImage(cvGetSize(img2Raw), IPL_DEPTH_8U, 1);

		ConvertToGray(img1Raw, img1);
		ConvertToGray(img2Raw, img2);

		cvGoodFeaturesToTrack(img1, eigImg, tmpImg, corners, cornerCount, lk_quality, lk_min_distance, null, 3, 0, 0.04);
		cvFindCornerSubPix(img1, corners, cornerCount.getValue(), lk_window_size.byValue(), lk_zero_zone.byValue(),
				lk_termination_criteria.byValue());

		// calling Lucas-Kanade algorithm
		byte[] featuresFound = new byte[MAX_CORNERS * 2];
		float[] featuresError = new float[MAX_CORNERS];

		CvSize pyramidSize = cvSize(img1.width + 8, img2.height/3);
		IplImage pyrA = cvCreateImage(pyramidSize.byValue(), IPL_DEPTH_32F, 1);
		IplImage pyrB = cvCreateImage(pyramidSize.byValue(), IPL_DEPTH_32F, 1);

		CvPoint2D32f[] corners2 = CvPoint2D32f.createArray(MAX_CORNERS);	// Note: has to be this! not new CvPoint2D32f[MAX_CORNERS]
		
		cvCalcOpticalFlowPyrLK(img1, img2, pyrA, pyrB, corners, corners2, cornerCount.getValue(), lk_window_size.byValue(),
				lk_pyr_level, featuresFound, featuresError, lk_termination_criteria.byValue(), 0);
		
		int cornersCountValue = cornerCount.getValue();
		
		float sumX = 0.0f;
		float sumY = 0.0f;
		int totalVectors = 0;
		allVectors = new ArrayList<Vector2D>(cornersCountValue);
		for (int i=0; i < cornersCountValue; i++)
		{
			if (featuresFound[i] == 0 || featuresError[i] > 550)
			{
				continue;
			}
			sumX += corners2[i].x - corners[i].x;
			sumY += corners2[i].y - corners[i].y;
			totalVectors++;
			
			allVectors.add(new Vector2D(corners[i].x, corners[i].y, corners2[i].x, corners2[i].y));
		}

		// Release IplImage objects
		cvReleaseImage(eigImg.pointerByReference());
		cvReleaseImage(tmpImg.pointerByReference());
		cvReleaseImage(img1.pointerByReference());
		cvReleaseImage(img2.pointerByReference());
		cvReleaseImage(pyrA.pointerByReference());
		cvReleaseImage(pyrB.pointerByReference());
		
		return new Vector2D(0f, 0f, sumX / totalVectors, sumY / totalVectors);
	}
	
		
	private static void ConvertToGray(IplImage original, IplImage dst)
	{
		if (original.nChannels > 1)
		{
			cvCvtColor(original, dst, CV_RGB2GRAY);
		}
		else
		{
			cvCopy(original, dst);
		}
	}
	
	public static Vector2D calculateAverage(List<Vector2D> vectors)
	{
		float sumX = 0;
		float sumY = 0;
		for (Vector2D vector : vectors)
		{
			sumX += (vector.getX2() - vector.getX1());
			sumY += (vector.getY2() - vector.getY1());
		}
		int vectorsLength = vectors.size();
		if (vectorsLength == 0) return new Vector2D(0f, 0f, 0f, 0f);
		
		return new Vector2D(0, 0, sumX / vectorsLength, sumY / vectorsLength);
	}
	
	public boolean isDebugDrawingEnabled() {
		return debugDrawingEnabled;
	}
	
	public void setDebugDrawingEnabled(boolean debugDrawingEnabled) {
		this.debugDrawingEnabled = debugDrawingEnabled;
	}

	public static double getLk_quality() {
		return lk_quality;
	}

	public static void setLk_quality(double lkQuality) {
		lk_quality = lkQuality;
	}

	public static double getLk_min_distance() {
		return lk_min_distance;
	}

	public static void setLk_min_distance(double lkMinDistance) {
		lk_min_distance = lkMinDistance;
	}

	public static CvSize getLk_zero_zone() {
		return lk_zero_zone;
	}

	public static void setLk_zero_zone(CvSize lkZeroZone) {
		lk_zero_zone = lkZeroZone;
	}

	public static CvSize getLk_window_size() {
		return lk_window_size;
	}

	public static void setLk_window_size(CvSize lkWindowSize) {
		lk_window_size = lkWindowSize;
	}

	public static CvTermCriteria getLk_termination_criteria() {
		return lk_termination_criteria;
	}

	public static void setLk_termination_criteria(
			CvTermCriteria lkTerminationCriteria) {
		lk_termination_criteria = lkTerminationCriteria;
	}

	public static int getLk_pyr_level() {
		return lk_pyr_level;
	}

	public static void setLk_pyr_level(int lkPyrLevel) {
		lk_pyr_level = lkPyrLevel;
	}
}
